
import glob
import torch
import datetime  # 用于获取当前时间，以便在日志中加上时间戳
import os
import h5py
import math
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import classification_report
import pandas as pd
import re
from tqdm import tqdm


class DataLabelsGenerate:
    def __init__(self):
        print("DataLabelsGenerate __init__")
        self.save_files_path = "F:\\EEGDataset\\MASS\\data_anno_savefiles"
        # self.save_files = glob.glob(os.path.join(self.save_files_path, "*.h5"))


    def generate_data_labels(self, file_names , batch_size):
        while True:
            for index, file_name in enumerate(file_names):

                with h5py.File(file_name, 'r') as f:
                    eeg_data_group = f['data']
                    channel_names = list(eeg_data_group.keys())  # 获取所有通道名称
                    data_list = []
                    # 遍历每个通道并加载数据
                    for channel in channel_names:
                        channel_data = eeg_data_group[channel][:]
                        data_tensor = torch.tensor(channel_data, dtype=torch.float32)
                        data_list.append(data_tensor.unsqueeze(1))
                    combined_data = torch.cat(data_list, dim=1)

                    labels = f['labels']['stage_labels'][:]
                    labels = torch.tensor(labels, dtype=torch.long)

                    indices = np.arange(len(combined_data))  # 生成索引
                    for i in range(0, len(combined_data), batch_size):
                        batch_indices = indices[i:i + batch_size]
                        batch_data = combined_data[batch_indices]

                        result = {
                            'batch_data': batch_data,
                            'labels': labels[batch_indices],
                            'file_name': file_name
                        }
                        yield result

